package edu.cn.nlsde.tmfst.mlda;

import common.voseAlias;
import edu.cn.nlsde.Utils.IOUtil;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.util.*;
import java.util.Map.Entry;

public class AliasSATM {

    public int P = 200;
    public int K = 20;

    public double alpha = 0.1;
    public double beta = 0.01;

    public double kAlpha;
    public double vBeta;

    public int MH_STEPS = 2;

    public int M;
    public int V;

    public int niters = 100;

    public int npw[][]; // P * V
    public int npwSum[]; // P * 1

    public int npk[][];
    public int npkSum[];

    public int nkw[][];
    public int nkwSum[];

    public List<voseAlias> q = new ArrayList<voseAlias>();

    public int d_assigns[][];
    public int z_assigns[][];

    public double eta[][];

    public double epsilon = 0.001;

    public double corpus_size = 0;

    public List<List<Integer>> docs = new ArrayList<List<Integer>>();
    public HashMap<String, Integer> w2i = new HashMap<String, Integer>();
    public HashMap<Integer, String> i2w = new HashMap<Integer, String>();

    public static void main(String args[]) {
        AliasSATM satm = new AliasSATM();
        satm.loadTxts("/Users/zuoyuan/Desktop/Experiment/20130601-origin");
        satm.init_model();
        satm.estimate();
        satm.printTopics(10);
    }

    public void loadTxts(String txtPath) {
        BufferedReader reader = IOUtil.getReader(txtPath);

        String line;
        try {
            line = reader.readLine();
            while (line != null) {
                List<Integer> doc = new ArrayList<Integer>();

                String[] tokens = line.trim().split("\\s+");
                for (String token : tokens) {
                    if (!w2i.containsKey(token)) {
                        w2i.put(token, w2i.size());
                        i2w.put(w2i.get(token), token);
                    }
                    doc.add(w2i.get(token));
                }
                docs.add(doc);
                line = reader.readLine();
            }
            reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        M = docs.size();
        V = w2i.size();
        return;
    }

    public void init_model() {
        kAlpha = K * alpha;
        vBeta = V * beta;

        eta = new double[M][P];

        npw = new int[V][P];
        npwSum = new int[P];

        npk = new int[P][K];
        npkSum = new int[P];

        nkw = new int[V][K];
        nkwSum = new int[K];

        z_assigns = new int[M][];
        d_assigns = new int[M][];

        for (int m = 0; m != M; m++) {
            int N = docs.get(m).size();
            corpus_size += N;

            z_assigns[m] = new int[N];
            d_assigns[m] = new int[N];
            int d = (int) Math.floor(Math.random() * P);
            for (int n = 0; n != N; n++) {
                int w = docs.get(m).get(n);
                int z = (int) Math.floor(Math.random() * K);

                npw[w][d]++;
                npwSum[d]++;

                npk[d][z]++;
                npkSum[d]++;

                nkw[w][z]++;
                nkwSum[z]++;

                z_assigns[m][n] = z;
                d_assigns[m][n] = d;
            }
        }

        this.update_eta();
        this.specific_init();
    }

    public void specific_init() {
        System.out.println("Initializing the alias tables ...");
        for (int w = 0; w != V; w++) {
            q.add(new voseAlias());
        }
        for (int w = 0; w < V; ++w) {
            q.get(w).init(P * K);
            generateQtable(w);
        }
    }

    public void generateQtable(int w) {
        q.get(w).wsum = 0.0;
        for (int p = 0; p != P; p++) {
            for (int k = 0; k != K; k++) {
                q.get(w).w[p * K + k] = (npk[p][k] + alpha) / (npkSum[p] + kAlpha)
                        * (nkw[w][k] + beta) / (nkwSum[k] + vBeta);
                q.get(w).wsum += q.get(w).w[p * K + k];
            }
        }
        q.get(w).constructTable();
    }

    public void update_eta() {
        eta = new double[M][P];

        for (int m = 0; m != M; m++) {
            double norm = 0.0;
            for (int p = 0; p != P; p++) {
                eta[m][p] = (npwSum[p] + 1.0) / (corpus_size + P);
                for (int n = 0; n != docs.get(m).size(); n++) {
                    int w = docs.get(m).get(n);
                    eta[m][p] *= (npw[w][p] + 1.0) / ((double) npwSum[p] + V);
                }
                norm += eta[m][p];
            }
            for (int p = 0; p != P; p++) {
                eta[m][p] /= norm;
            }
        }
    }

    public void sample_doc(int m) {
        int N = docs.get(m).size();

        List<Integer> dList = new ArrayList<Integer>();
        for (int p = 0; p != P; p++) {
            if (eta[m][p] > epsilon)
                dList.add(p);
        }

        for (int n = 0; n != N; n++) {
            int w = docs.get(m).get(n);

            int d = d_assigns[m][n];
            int z = z_assigns[m][n];
            int new_d = -1;
            int new_z = -1;

            double[] p = new double[dList.size() * K];

            nkw[w][z]--;
            nkwSum[z]--;

            npk[d][z]--;
            npkSum[d]--;

            npw[w][d]--;
            npwSum[d]--;

            double psum = 0.0, qsum = 0.0;
            int index = 0;
            for (int pd : dList) {
                for (int k = 0; k != K; k++) {
                    psum += eta[m][pd] * (npk[pd][k] + alpha) / (npkSum[pd] + kAlpha) *
                            (nkw[w][k] + beta) / (nkwSum[k] + vBeta);
                    qsum += eta[m][pd] * q.get(w).wsum;
                    p[index] = psum;
                    index++;
                }
            }

            double select_pr = psum / qsum;

            for (int s = 0; s != MH_STEPS; s++) {
                if (Math.random() < select_pr) {
                    double u = Math.random() * psum;
                    for (int i = 0; i != index; i++) {
                        if (u < p[i]) {
                            new_d = dList.get(i / K);
                            new_z = i % K;
                            break;
                        }
                    }
                } else {
                    q.get(w).noSamples++;
                    if (q.get(w).noSamples > (P * K / (double) MH_STEPS)) {
                        generateQtable(w);
                    }
                    int temp = q.get(w).sample(Math.random(), Math.random());
                    new_d = temp / K;
                    new_z = temp % K;
                }
            }

            if (z != new_z || d != new_d) {
                //2. Find acceptance probability
                double temp_old = (npk[d][z] + alpha) / (npkSum[d] + kAlpha)
                        * (nkw[w][z] + beta) / (nkwSum[z] + vBeta);
                double temp_new = (npk[new_d][new_z] + alpha) / (npkSum[new_d] + kAlpha)
                        * (nkw[w][new_z] + beta) / (nkwSum[new_z] + vBeta);
//				double acceptance = eta[m][new_d] / eta[m][d]
//						* temp_new / temp_old
//						*(eta[m][d] * q.get(w).w[d*K+z]) /
//						(eta[m][new_d] * q.get(w).w[new_d*K+new_z]);

                double acceptance = temp_new / temp_old
                        * q.get(w).w[d * K + z] / q.get(w).w[new_d * K + new_z];

                if (Math.random() < acceptance) {
                    d = new_d;
                    z = new_z;
                }
            }

            nkw[w][z]++;
            nkwSum[z]++;

            npk[d][z]++;
            npkSum[d]++;

            npw[w][d]++;
            npwSum[d]++;

            d_assigns[m][n] = d;
            z_assigns[m][n] = z;
        }
    }

    public void sample_d_z(int m, int n) {
        int w = docs.get(m).get(n);
        int d = d_assigns[m][n];
        int z = z_assigns[m][n];

        npk[d][z]--;
        npkSum[d]--;
        nkw[w][z]--;
        nkwSum[z]--;

        npw[w][d]--;
        npwSum[d]--;

        double pTable[][] = new double[P][K];
        double pSum = 0.0;
        for (int p = 0; p != P; p++) {
            if (eta[m][p] < epsilon)
                continue;
            for (int k = 0; k != K; k++) {
                pTable[p][k] = eta[m][p] * (npk[p][k] + alpha) / (npkSum[p] + K * alpha)
                        * (nkw[w][k] + beta) / (nkwSum[k] + V * beta);
                pSum += pTable[p][k];
            }
        }

        double r = Math.random() * pSum;
        pSum = 0.0;
        boolean brk = false;
        for (int p = 0; p != P; p++) {
            if (brk == true)
                break;
            if (eta[m][p] < epsilon)
                continue;
            for (int k = 0; k != K; k++) {
                pSum += pTable[p][k];
                if (r < pSum) {
                    d = p;
                    z = k;
                    brk = true;
                    break;
                }
            }
        }

        npk[d][z]++;
        npkSum[d]++;
        nkw[w][z]++;
        nkwSum[z]++;

        npw[w][d]++;
        npwSum[d]++;

        z_assigns[m][n] = z;
        d_assigns[m][n] = d;
    }

    public void estimate() {
        for (int iter = 0; iter != niters; iter++) {
            System.out.println("Alias SATM Iteration " + iter + " ...");
            long begin = System.currentTimeMillis();
            for (int m = 0; m != M; m++) {
//				for (int n = 0; n != docs.get(m).size(); n++) {
//					sample_d_z(m, n);
//				}
                sample_doc(m);
            }
            this.update_eta();
            System.out.println((System.currentTimeMillis() - begin) / 1000.0);
        }
    }

    public double[][] computePhi() {
        double[][] phi = new double[K][V];
        for (int k = 0; k != K; k++) {
            for (int v = 0; v != V; v++) {
                phi[k][v] = (nkw[v][k] + beta) / (nkwSum[k] + V * beta);
            }
        }
        return phi;
    }

    public ArrayList<List<Entry<String, Double>>> sortedTopicWords(
            double[][] phi, int T) {
        ArrayList<List<Entry<String, Double>>> res = new ArrayList<List<Entry<String, Double>>>();
        for (int k = 0; k != T; k++) {
            HashMap<String, Double> term2weight = new HashMap<String, Double>();
            for (String term : w2i.keySet())
                term2weight.put(term, phi[k][w2i.get(term)]);

            List<Entry<String, Double>> pairs = new ArrayList<Entry<String, Double>>(
                    term2weight.entrySet());
            Collections.sort(pairs, new Comparator<Entry<String, Double>>() {
                public int compare(Entry<String, Double> o1,
                                   Entry<String, Double> o2) {
                    return (o2.getValue().compareTo(o1.getValue()));
                }
            });
            res.add(pairs);
        }
        return res;
    }

    public void printTopics(int top_n) {
        double[][] phi = computePhi();
        ArrayList<List<Entry<String, Double>>> pairsList = this
                .sortedTopicWords(phi, K);
        for (int k = 0; k != K; k++) {
            System.out.println("Topic " + k + ":");
            for (int i = 0; i != top_n; i++) {
                System.out.println(pairsList.get(k).get(i).getKey() + " "
                        + pairsList.get(k).get(i).getValue());
            }
        }
    }

    public void savePhi(String path) {
        BufferedWriter writer = IOUtil.getWriter(path);

        double[][] phi = computePhi();
        int K = phi.length;
        assert K > 0;
        int V = phi[0].length;

        try {
            for (int k = 0; k != K; k++) {
                for (int v = 0; v != V; v++)
                    writer.append(phi[k][v] + " ");
                writer.append("\n");
            }
            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void saveWordmap(String path) {
        BufferedWriter writer = IOUtil.getWriter(path);

        try {
            for (String word : w2i.keySet())
                writer.append(word + "\t" + w2i.get(word) + "\n");

            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

}
